# coding:utf-8
'''
多层+dropout的网络
'''
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import numpy as np
# 初始化权重和偏置项
def init_weights_and_bias(shapeW,shapeB):
    return tf.Variable(tf.truncated_normal(shape=shapeW,stddev=0.01),dtype='float'),tf.Variable(np.zeros(shapeB),dtype='float')
# 网络结构
def model(x,w_1,b_1,w_2,b_2,w_3,b_3,keep_prob):
    h1=tf.nn.relu(tf.matmul(x,w_1)+b_1)
    h2=tf.nn.relu(tf.matmul(h1,w_2)+b_2)
    h2=tf.nn.dropout(h2,keep_prob=keep_prob) # 加dropout防止过拟合
    return tf.matmul(h2,w_3)+b_3

# 提供占位符
x_input = tf.placeholder(dtype='float',shape=[None,784])
y_input = tf.placeholder(dtype='float',shape=[None,10])
keep_prob = tf.placeholder(dtype='float')

# 载入数据
mnist = input_data.read_data_sets('MNIST_data/',one_hot=True)
x_train,y_train,x_test,y_test = mnist.train.images,mnist.train.labels,mnist.test.images,mnist.test.labels

#生成权重和偏执项
w_1,b_1= init_weights_and_bias([784,625],[625])
w_2,b_2 = init_weights_and_bias([625,625],[625])
w_3,b_3 =init_weights_and_bias([625,10],[10])
# 输出学习值
y_model = model(x_input,w_1,b_1,w_2,b_2,w_3,b_3,keep_prob)
# 设计损失函数
cross_entory = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=y_model,labels=y_input))
# 选择优化器
train_op = tf.train.RMSPropOptimizer(0.001,0.9).minimize(cross_entory)
# 预测结果
predict_op = tf.argmax(y_model,1)

# 创建会话
with tf.Session() as sess:
    # 初始化所有变量
    tf.global_variables_initializer().run()
    epoch_nums = 50 # 迭代总轮次
    batch_size = 128 # 每次batch的数据量
    for i in range(epoch_nums):
        for start,end in zip(range(0,len(x_train),batch_size),range(batch_size,len(x_train),batch_size)):
            sess.run(train_op,feed_dict={x_input:x_train,y_input:y_train,keep_prob:0.5})
        train_acc = np.mean(np.argmax(y_train,1) == sess.run(predict_op,feed_dict={x_input:x_train,keep_prob:1.0}))
        test_acc = np.mean(np.argmax(y_test,1) == sess.run(predict_op,feed_dict={x_input:x_test,keep_prob:1.0}))
        print('epoch:',i+1,'train_acc:',train_acc*100,'%','test_acc:',test_acc*100,'%')